# -*- coding:utf-8 -*-
################################################################################
# This script for select gpu
#
# Written by CathyQ
################################################################################

import os


def gpu_select(gpu_count):
    """
    Select the top k free memory GPU for working

    Args:
        gpu_count: int - number of gpu for selecting
    """
    os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
    memory_gpu = [int(x.split()[2]) for x in open('tmp', 'r').readlines()]
    print("GPUs free memory: {}".format(memory_gpu))
    gpu = [(item, i) for i, item in enumerate(memory_gpu)]
    gpu.sort(key=lambda x: x[0])
    gpu_selected = ([str(item[1]) for item in gpu])[-gpu_count:]
    print("GPUs : {} has been selected.".format(gpu_selected))
    os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(gpu_selected)
    os.system('rm tmp')
    return True
